# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import collections
import json

import cc.coin1.prophet.py.prophet.op_registry as native_registry
from experimental.prophet.graph import graph_def_pb2
from experimental.prophet.graph.shape import Shape
from experimental.prophet.graph.util import from_data_type_proto, to_data_type_proto
from experimental.prophet.kernels.op_kernel import _ArgDef, _VariadicInput
from experimental.prophet.kernels.op_kernel_context import PyOpKernelContext

_ArgRef = collections.namedtuple('ArgRef',
                                 ['node_name', 'output_index', 'input_name', 'dtype', 'shape'])


def _to_arg_def(arg):
  if isinstance(arg, _ArgDef):
    return arg

  elif arg in (float, int, bool):
    return ArgDef_(None, arg, Shape(), None)

  elif isinstance(arg, (list, tuple)):
    if len(arg) == 3:
      return ArgDef_(arg[0], arg[1], arg[2], None)
    elif len(arg) == 4:
      return ArgDef_(*args)

  raise ValueError(arg)


def _to_arg_def_list(args):
  assert isinstance(args, list)

  ret = []
  for idx, arg in enumerate(args):
    arg = _to_arg_def(arg)
    if arg.name is None:
      arg = _ArgDef(str(idx), *arg[1:])
    ret.append(arg)
  return ret


def _convert_arg_def_to_proto(arg):
  arg_proto = graph_def_pb2.ArgDefProto()
  arg_proto.name = arg.name
  arg_proto.dtype = to_data_type_proto(arg.dtype)
  arg_proto.shape.CopyFrom(Shape(arg.shape).to_proto())
  arg_proto.attrs_json = json.dumps(arg.attrs) if arg.attrs is not None else ''
  return arg_proto


def _get_op_def_proto_from_kernel_cls(kernel_cls):
  op_def = graph_def_pb2.OpDefProto()

  if not kernel_cls.name:
    raise ValueError('invalid name')
  op_def.name = kernel_cls.name

  op_def.deprecated = kernel_cls.deprecated
  op_def.is_stateful = kernel_cls.is_stateful
  op_def.support_pickling = kernel_cls.support_pickling
  op_def.allow_null_input = kernel_cls.allow_null_input
  op_def.is_delayed_op = kernel_cls.is_delayed_op
  return op_def


class PyOp:
  def __init__(self, kernel_cls):
    self._kernel_cls = kernel_cls
    self.op_def = _get_op_def_proto_from_kernel_cls(kernel_cls)

  def instantiate(self, binding):
    ctx = PyOpKernelContext(self._kernel_cls)
    binding.instance = ctx
    binding.init_func = ctx.init
    binding.compute_func = ctx.compute
    binding.on_null_input_func = ctx.on_null_input

  def get_input_args(self, serialized_req):
    req = graph_def_pb2.OpInputDefRequestProto()
    req.ParseFromString(serialized_req)

    attrs = None
    if req.attrs_json:
      attrs = json.loads(req.attrs_json)
    inputs = self._kernel_cls.get_input_args(attrs)
    if inputs is None:
      raise ValueError('input is not defined')

    res = graph_def_pb2.OpInputDefResponseProto()
    if inputs == _VariadicInput:
      res.variadic_input = True
    else:
      res.input_args.extend([_convert_arg_def_to_proto(arg) for arg in _to_arg_def_list(inputs)])
    return res.SerializeToString()

  def get_output_args(self, serialized_req):
    req = graph_def_pb2.OpOutputDefRequestProto()
    req.ParseFromString(serialized_req)

    attrs = None
    if req.attrs_json:
      attrs = json.loads(req.attrs_json)

    inputs = []
    for arg_proto in req.inputs:
      arg = _ArgRef(node_name=arg_proto.node_name,
                    output_index=arg_proto.output_index,
                    input_name=arg_proto.input_name,
                    dtype=from_data_type_proto(arg_proto.dtype),
                    shape=Shape.from_proto(arg_proto.shape))
      inputs.append(arg)
    outputs = self._kernel_cls.get_output_args(tuple(inputs), attrs)
    if outputs is None:
      raise ValueError('output is not defined')

    res = graph_def_pb2.OpOutputDefResponseProto()
    res.output_args.extend([_convert_arg_def_to_proto(arg) for arg in _to_arg_def_list(outputs)])
    return res.SerializeToString()


def register_op(kernel_cls):
  op = PyOp(kernel_cls)
  native_registry.register_py_op(op.op_def.SerializeToString(),
                                 op.instantiate,
                                 op.get_input_args,
                                 op.get_output_args)


def list_ops():
  return native_registry.list_ops()


def has_op(name):
  return native_registry.has_op(name)


def get_input_args(name, attrs, node_name=None):
  req = graph_def_pb2.OpInputDefRequestProto()
  if node_name is not None:
    req.name = node_name
  req.op = name
  if attrs is not None:
    req.attrs_json = json.dumps(attrs)
  res_serialized = native_registry.get_input_args(req.SerializeToString())
  res = graph_def_pb2.OpInputDefResponseProto()
  res.ParseFromString(res_serialized)
  return res


def get_output_args(name, attrs, inputs, node_name=None):
  req = graph_def_pb2.OpOutputDefRequestProto()
  if node_name is not None:
    req.name = node_name
  req.op = name
  if attrs is not None:
    req.attrs_json = json.dumps(attrs)

  if isinstance(inputs, collections.Mapping):
    inputs = list(inputs.items())
  else:
    inputs = [(None, input) for input in inputs]

  for input_name, input_var in inputs:
    arg_ref = req.inputs.add()
    arg_ref.node_name = input_var.node.name
    arg_ref.output_index = input_var.output_idx
    arg_ref.input_name = input_name
    arg_ref.dtype = to_data_type_proto(input_var.dtype)
    arg_ref.shape.CopyFrom(Shape(input_var.shape).to_proto())

  res_serialized = native_registry.get_output_args(req.SerializeToString())
  res = graph_def_pb2.OpOutputDefResponseProto()
  res.ParseFromString(res_serialized)
  return res


def get_op_def(name):
  res_serialized = native_registry.get_op_def(name)
  res = graph_def_pb2.OpDefProto()
  res.ParseFromString(res_serialized)
  return res
